package com.benromberg.jameshook.internal;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

import javassist.ByteArrayClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;

import org.apache.log4j.Logger;

import com.benromberg.jameshook.ClassHook;
import com.benromberg.jameshook.HookMethod.HookMode;
import com.benromberg.jameshook.JamesHook;

public class Transformer implements ClassFileTransformer {

	private Logger logger = Logger.getLogger(getClass());
	private ClassPool pool = ClassPool.getDefault();

	private static Set<String> loadedClasses = new HashSet<String>();

	public static Set<String> getLoadedClasses() {
		return loadedClasses;
	}

	@Override
	public byte[] transform(ClassLoader loader, String classNameWithSlashes, Class<?> classBeingRedefined,
			ProtectionDomain protectionDomain, byte[] classfileBuffer) {
		String className = classNameWithSlashes.replace('/', '.');
		loadedClasses.add(className);
		logger.debug("checking class " + className);
		if (!JamesHook.getHooks().containsKey(className)) {
			return null;
		}
		logger.info("found hook for class " + className);
		ClassHook<?> hook = JamesHook.getHooks().get(className);
		try {
			if (loader != null) {
				pool.insertClassPath(new LoaderClassPath(loader));
			}
			if (classfileBuffer != null) {
				pool.insertClassPath(new ByteArrayClassPath(className, classfileBuffer));
			}

			CtClass targetClass = pool.get(className);

			for (CtMethod method : targetClass.getDeclaredMethods()) {
				String runtimeInfo = "new com.benromberg.jameshook.RuntimeInfoImpl(\"" + method.getName() + "\")";
				String callString = "((com.benromberg.jameshook.ClassHook) com.benromberg.jameshook.JamesHook#getHooks().get(\""
						+ className + "\")).run(this, " + runtimeInfo + ");";
				if (hook.getHookMode() == HookMode.BEFORE_AND_AFTER || hook.getHookMode() == HookMode.BEFORE) {
					logger.info("inserting hook before method " + method.getLongName());
					method.insertBefore(callString);
				} else if (hook.getHookMode() == HookMode.BEFORE_AND_AFTER || hook.getHookMode() == HookMode.AFTER) {
					logger.info("inserting hook after method " + method.getLongName());
					method.insertAfter(callString);
				} else if (hook.getHookMode() == HookMode.REPLACE) {
					logger.info("inserting hook replacing method " + method.getLongName());
					method.setBody(callString);
				}
			}
			return targetClass.toBytecode();
		} catch (Exception e) {
			logger.error("transformation exception", e);
			return null;
		}
	}

}
